

import onnx

# 加载模型
model = onnx.load("end2end.onnx")
graph = model.graph

# 1. 查找并移除ArgMax节点
argmax_node = next(n for n in graph.node if n.op_type == "ArgMax")
graph.node.remove(argmax_node)

# 2. 获取ArgMax的输入名称（即上游节点的输出）
input_to_argmax = argmax_node.input[0]

# 3. 查找产生该输入的上游节点
prev_node = next(n for n in graph.node if input_to_argmax in n.output)

# 4. 重定向上游节点的输出到原始输出名称
original_output_name = graph.output[0].name
prev_node.output[0] = original_output_name  # 替换为模型原始输出名

# 5. 更新输出张量的数据类型和形状
output_tensor = graph.output[0]

# 修改数据类型为float32（对应ONNX类型1）
output_tensor.type.tensor_type.elem_type = 1

# 获取类别数（假设您知道实际值，例如19）
num_classes = 19  # 替换为您的实际类别数

# 更新输出形状为 [1, num_classes, height, width]
shape = output_tensor.type.tensor_type.shape

# 确保形状有4个维度（batch, channels, height, width）
if len(shape.dim) != 4:
    # 如果原始形状不是4D，重建维度
    shape.dim.clear()
    shape.dim.add().dim_value = 1              # batch=1
    shape.dim.add().dim_value = num_classes     # channels=类别数
    # 高度和宽度保持动态（从上游节点推断）
    shape.dim.add().dim_param = 'height'        # 动态高度
    shape.dim.add().dim_param = 'width'         # 动态宽度
else:
    # 直接修改通道维度
    shape.dim[1].dim_value = num_classes

# 6. 清理不再使用的ValueInfo（可选）
for i in range(len(graph.value_info)):
    if graph.value_info[i].name == input_to_argmax:
        del graph.value_info[i]
        break

# 验证并保存模型
onnx.checker.check_model(model)
onnx.save(model, "modified_model.onnx")



